from fastapi import APIRouter, File, UploadFile, BackgroundTasks
from loguru import logger
from ultralytics import YOLO

from app.dependencies.model_deps import ModelDep
from app.schema.response import BaseResponse
from app.schema.detection import DetectionResponse, Base64Request
from app.services import detection, image_processing
from app.utils import base64

router = APIRouter()


@router.post("/detect/file", response_model=DetectionResponse)
async def detect_from_file(
        file: UploadFile = File(...),
        model: YOLO = ModelDep,
        background_tasks: BackgroundTasks = BackgroundTasks()
):
    try:
        # 验证并读取文件
        content = await file.read()
        image_processing.validate_image_size(content)
        image = image_processing.bytes_to_cvimage(content)

        # 异步执行检测
        detections = await detection.predict_image(model, image)
        logger.info("----------" + detections)
        # 清理内存
        background_tasks.add_task(file.file.close)
        return BaseResponse.success(data=detections)
    except Exception as e:
        return BaseResponse.error(code=1, message=str(e))


@router.post("/detect/base64", response_model=DetectionResponse)
async def detect_from_base64(
        request: Base64Request,
        model: YOLO = ModelDep
):
    try:
        # 转换Base64
        content = base64.base64_to_bytes(request.image_base64)
        image_processing.validate_image_size(content)
        image = image_processing.bytes_to_cvimage(content)

        # 执行检测
        detections = await detection.predict_image(model, image)

        return BaseResponse.success(data=detections)
    except Exception as e:
        return BaseResponse.error(code=1, message=str(e))